import logging

from flask import Blueprint
from flask import Response
from flask import jsonify
from flask import request
from sqlalchemy.orm.session import Session

from airflow.patsnap.service import task_service
from airflow.models.idata_task import Task
from airflow.models.idata_task_alert import TaskAlert
from airflow.models.idata_workflow import WorkFlow
from airflow.models.taskinstance import TaskInstance
from airflow.utils import dates
from airflow.utils.session import provide_session
from airflow.utils.state import State
from airflow.www.api import exception_handle

log = logging.getLogger(__name__)

task_bp = Blueprint('task_bp', __name__)


@task_bp.route('/health', methods=['GET'])
@exception_handle
def health():
    return jsonify(code=0, message='ok', data=None)


@task_bp.route('/tasks/<int:task_id>/trigger', methods=['POST'])
@exception_handle
@provide_session
def task_run(task_id, session: Session = None):
    task = Task.get_task(task_id, session=session)
    if task:
        if task.task_concurrency:
            running = session.query(TaskInstance).filter(TaskInstance.task_id == 'task-{}'.format(task_id),
                                                         TaskInstance.state == State.RUNNING).all()
            if len(running) >= task.task_concurrency:
                return jsonify(code=1, message='task_concurrency exceed {}'.format(task.task_concurrency),
                               data=None)
        resp = task_service.run_by_scheduler(task)
        return jsonify(code=0, message='ok', data=resp)
    else:
        return jsonify(code=404, message='task not found', data=None)


@task_bp.route('/tasks/<int:task_id>', methods=['GET'])
@exception_handle
def task_get(task_id):
    task = Task.get_task(task_id)
    if task:
        return jsonify(code=0, message='ok', data=task.to_json())
    else:
        return jsonify(code=404, message='task not found', data=None)


@task_bp.route('/tasks', methods=['POST'])
@exception_handle
def task_insert():
    data = request.get_json()
    task = Task(**data)
    task_service.insert_task(task)
    return jsonify(code=0, message='ok', data=task.to_json())


@task_bp.route('/tasks/file', methods=['POST'])
@exception_handle
def task_file_add():
    data = request.get_json()
    task_id = data.get('task_id', None)
    file_path = data.get('file_path', None)
    if task_id and file_path:
        task_service.add_task_file(task_id, file_path)
    else:
        return jsonify(code=400, message='task_id or file_path is null', data=None)
    return jsonify(code=0, message='ok', data=None)


@task_bp.route('/tasks/file', methods=['DELETE'])
@exception_handle
def task_file_delete():
    data = request.get_json()
    task_id = data.get('task_id', None)
    file_path = data.get('file_path', None)
    if task_id and file_path:
        task_service.delete_task_file(task_id, file_path)
    else:
        return jsonify(code=400, message='task_id or file_path is null', data=None)
    return jsonify(code=0, message='ok', data=None)


@task_bp.route('/tasks', methods=['GET'])
@exception_handle
def task_query():
    project_id = request.args.get('project_id', None)
    folder_id = request.args.get('folder_id', None)
    if project_id is None and folder_id is None:
        return jsonify(code=1, message='project_id or folder_id 不能为空', data=None)
    name = request.args.get('name', None)
    tasks = Task.query_task(project_id=project_id, folder_id=folder_id, name=name)
    rs = [task.to_json() for task in tasks]
    return jsonify(code=0, message='ok', data=rs)


@task_bp.route('/tasks/<int:task_id>', methods=['PUT'])
@exception_handle
def task_update(task_id):
    data = request.get_json()
    if 'id' not in data:
        data['id'] = task_id
    task = Task(**data)
    task_service.update_task(task)
    return jsonify(code=0, message='ok', data=task.to_json())


@task_bp.route('/tasks/<int:task_id>/command', methods=['GET'])
@exception_handle
def task_cmd(task_id):
    task = Task.get_task(task_id)
    if task:
        command = task.get_task_command()
        return jsonify(code=0, message='ok', data=command)
    else:
        return jsonify(code=404, message='task not found', data=None)


@task_bp.route('/tasks/<int:task_id>/workflow', methods=['GET'])
@exception_handle
def task_workflow(task_id):
    rs = [graph.to_json(include_content=False) for graph in WorkFlow.query_task_workflow(task_id)]
    return jsonify(code=0, message='ok', data=rs)


@task_bp.route('/tasks/instance', methods=['GET'])
@exception_handle
@provide_session
def task_instance_query(size=20, session: Session = None):
    dag_id = request.args.get('dag_id', None)
    task_id = request.args.get('task_id', None)
    execution_date = request.args.get('execution_date', None)
    if dag_id or task_id:
        qry = session.query(TaskInstance)
        if dag_id:
            qry = qry.filter(TaskInstance.dag_id == dag_id)
        if task_id:
            qry = qry.filter(TaskInstance.task_id == task_id)
        if execution_date:
            assert len(str(execution_date)) == 13
            date = dates.from_milliseconds(execution_date)
            qry = qry.filter(TaskInstance.execution_date == date)
        rs = [ti.to_json() for ti in qry.order_by(TaskInstance.execution_date.desc()).limit(size).all()]
        return jsonify(code=0, message='ok', data=rs)
    else:
        return jsonify(code=400, message='dag_id and task_id is null', data=None)


@task_bp.route('/tasks/kill', methods=['POST'])
@exception_handle
@provide_session
def task_instance_kill(session: Session = None):
    data = request.get_json()
    dag_id = data.get('dag_id', None)
    task_id = data.get('task_id', None)
    execution_date = data.get('execution_date', None)
    if dag_id and task_id and execution_date:
        assert len(str(execution_date)) == 13
        date = dates.from_milliseconds(execution_date)
        ti = session.query(TaskInstance).filter(TaskInstance.dag_id == dag_id,
                                                TaskInstance.task_id == task_id,
                                                TaskInstance.execution_date == date).first()
        if ti is not None:
            ti.set_state(State.SHUTDOWN)
            return jsonify(code=0, message='ok', data=ti.to_json())
        else:
            return jsonify(code=404, message='task instance not found', data=None)
    else:
        return jsonify(code=400, message='dag_id or task_id or execution_date is null', data=None)


@task_bp.route('/tasks/alert', methods=['POST'])
@exception_handle
def insert_task_alert():
    data = request.get_json()
    task_id = data.get('task_id', None)
    alert_id = data.get('alert_id', None)
    if task_id and alert_id:
        TaskAlert(task_id=task_id, alert_id=alert_id).insert()
    else:
        return jsonify(code=400, message='task_id or alert_id is None', data=None)


@task_bp.route('/tasks/alert', methods=['DELETE'])
@exception_handle
def delete_task_alert():
    data = request.get_json()
    task_id = data.get('task_id', None)
    alert_id = data.get('alert_id', None)
    if task_id and alert_id:
        TaskAlert(task_id=task_id, alert_id=alert_id).delete()
    else:
        return jsonify(code=400, message='task_id or alert_id is None', data=None)


@task_bp.route('/tasks/log', methods=['GET'])
@exception_handle
def get_task_log():
    hostname = request.args.get('hostname', None)
    if hostname:
        rs = task_service.get_task_log(hostname)
        return jsonify(code=0, message='ok', data=rs)
    else:
        return jsonify(code=400, message='hostname is None', data=None)


@task_bp.route('/tasks/log/download', methods=['GET'])
@exception_handle
def download_task_log():
    hostname = request.args.get('hostname', None)
    if hostname:
        return Response(task_service.download_task_log(hostname), content_type='application/octet-stream')
    else:
        return jsonify(code=400, message='hostname is None', data=None)


@task_bp.route('/tasks/<int:task_id>/blood', methods=['POST'])
@exception_handle
def insert_task_blood(task_id):
    blood = request.get_json()
    Task.update_blood(task_id, blood)
    return jsonify(code=0, message='ok', data=None)
